import os
from PIL import Image
import concurrent.futures
import argparse
import logging
import threading
import time

# 设置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

# 创建一个 Event 对象，用于在找到目标图片后通知所有线程停止
stop_event = threading.Event()
found_image_path = None
lock = threading.Lock()

def check_image_size(image_path):
    """
    检查单个图片的尺寸。如果尺寸小于256x256，则设置停止事件。
    """
    # 如果已经找到，则快速退出
    if stop_event.is_set():
        return

    global found_image_path
    try:
        with Image.open(image_path) as img:
            width, height = img.size
            if width < 256 or height < 256:
                # 使用锁来确保线程安全地更新共享变量和设置事件
                with lock:
                    if not found_image_path: # 确保只记录第一个找到的
                        found_image_path = image_path
                        logging.info(f"Found image with size smaller than 256px: '{image_path}' ({width}x{height})")
                        stop_event.set() # 设置事件，通知其他线程停止
    except IOError:
        # 忽略无法打开的文件
        pass
    except Exception as e:
        logging.warning(f"Could not process file '{os.path.basename(image_path)}': {e}")


def main():
    parser = argparse.ArgumentParser(description="Check for images smaller than 256px in a directory using multiple threads.")
    parser.add_argument("src_dir", help="The source directory to scan for images.")
    parser.add_argument("--workers", type=int, default=16, help="Number of worker threads (default: number of CPUs).")
    args = parser.parse_args()

    src_dir = args.src_dir
    if not os.path.isdir(src_dir):
        logging.error(f"Source directory not found: {src_dir}")
        return

    # 获取所有文件的路径
    all_files = [os.path.join(src_dir, f) for f in os.listdir(src_dir) if os.path.isfile(os.path.join(src_dir, f))]
    
    if not all_files:
        logging.info("No files found in the directory.")
        return

    logging.info(f"Scanning {len(all_files)} files in '{src_dir}'...")
    start_time = time.time()

    # 使用线程池来并发检查图片
    with concurrent.futures.ThreadPoolExecutor(max_workers=args.workers) as executor:
        # 提交所有任务
        futures = [executor.submit(check_image_size, file_path) for file_path in all_files]
        
        # 等待直到 stop_event 被设置
        stop_event.wait()

        # 事件被设置后，取消所有仍在运行的未来任务
        logging.info("Stopping all threads as a small image has been found.")
        for future in futures:
            if not future.done():
                future.cancel()
        
        # 确保线程池在退出 with 块之前关闭
        executor.shutdown(wait=False, cancel_futures=True)


    end_time = time.time()
    logging.info(f"Scan finished in {end_time - start_time:.2f} seconds.")

    if found_image_path:
        print(f"\nResult: Yes, found at least one image smaller than 256px.")
        print(f"First occurrence: {found_image_path}")
    else:
        # 如果循环正常结束（没有通过 stop_event 提前退出）
        print("\nResult: No, all images are 256px or larger.")


if __name__ == "__main__":
    main()
